package timingwheel

import (
	"sync"
	"sync/atomic"
	"time"
)

// MultiLevelTimeWheel 定义多级时间轮的数据结构
type MultiLevelTimeWheel struct {
	wheels       []*SingleTimeWheel // 时间轮层级（0为最底层）
	tickDuration time.Duration      // 时间轮的嘀嗒间隔
	stopChan     chan struct{}      // 停止信号
	taskChan     chan *Task         // 任务队列
	taskMap      sync.Map           // 任务映射表
	tidGenerator atomic.Int64       // 任务ID生成器
	isRunning    atomic.Bool        // 运行状态
	waitGroup    sync.WaitGroup     // 等待组
}

// NewMultiLevelTimeWheel 新建一个多级时间轮
func NewMultiLevelTimeWheel(tickDuration time.Duration, slotNums []int) *MultiLevelTimeWheel {
	if tickDuration <= 0 {
		panic("timewheel: tick duration must be greater than zero")
	}
	if len(slotNums) == 0 {
		panic("timewheel: at least one time wheel level is required")
	}

	wheels := &MultiLevelTimeWheel{
		wheels:       make([]*SingleTimeWheel, len(slotNums)),
		tickDuration: tickDuration,
		stopChan:     make(chan struct{}),
		taskChan:     make(chan *Task, 1024),
	}

	// 从最底层开始逐级创建时间轮
	for i, slotNum := range slotNums {
		if i == 0 {
			wheels.wheels[i] = NewSingleTimeWheel(tickDuration, slotNum)
			continue
		}

		// 上一级时间轮转动一圈的时间作为当前级的滴答时间
		tickDuration = time.Duration(slotNums[i-1]) * tickDuration
		wheels.wheels[i] = NewSingleTimeWheel(tickDuration, slotNum)
	}

	return wheels
}

// Start 启动时间轮
func (own *MultiLevelTimeWheel) Start() {
	if own.isRunning.Swap(true) {
		return
	}

	// 启动各级时间轮
	for _, wheel := range own.wheels {
		wheel.Start()
	}

	own.waitGroup.Add(1)
	go own.run()
}

func (own *MultiLevelTimeWheel) run() {
	own.waitGroup.Done()

	for {
		select {
		case task := <-own.taskChan:
			own.addTaskToWheel(task)
		case <-own.stopChan:
			for _, wheel := range own.wheels {
				wheel.Stop()
			}
			return
		}
	}
}

// Stop 停止时间轮
func (own *MultiLevelTimeWheel) Stop() {
	if !own.isRunning.Load() {
		return
	}

	close(own.stopChan)
	own.waitGroup.Wait()
	own.isRunning.Store(false)
}

// AddTask 添加任务
func (own *MultiLevelTimeWheel) AddTask(delay time.Time, job func(), period time.Duration, repeats ...int) int64 {
	task := &Task{
		id:      own.tidGenerator.Add(1),
		job:     job,
		delay:   delay.UnixNano(),
		repeats: 1,
	}

	if period > 0 {
		task.period = period
	}

	if task.period > 0 && len(repeats) > 0 {
		if repeats[0] < 0 {
			task.repeats = repeats[0]
		} else {
			task.repeats += repeats[0]
		}
	}

	own.taskChan <- task
	return task.id
}

// addTaskToWheel 内部添加任务方法（自动选择合适的时间轮层级）
func (own *MultiLevelTimeWheel) addTaskToWheel(task *Task) {
	// 检查是否是否还需要执行
	if task.repeats == 0 {
		return
	}

	// 计算应该放入哪一级时间轮
	exp := task.delay - time.Now().UnixNano()
	base := int64(1)
	var wheel *SingleTimeWheel
	for _, wheel = range own.wheels {
		// 计算当前级时间轮的最大表示范围
		wheelTickDuration := int64(wheel.slotNumber)
		if exp < wheelTickDuration*base {
			break
		}

		base *= wheelTickDuration
	}

	own.addTaskToSlot(wheel, task)
}

// addTaskToSlot 内部添加任务方法
func (own *MultiLevelTimeWheel) addTaskToSlot(wheel *SingleTimeWheel, task *Task) {
	// 只要 task.repeats == 0，则表明任务执行不需要在执行
	if task.repeats == 0 {
		return
	}

	delay := task.delay - time.Now().UnixNano()
	if delay <= int64(own.tickDuration) {
		own.callback(task)
		return
	}

	// 计算任务应该在多少个时间滴答后执行
	ticks := delay / int64(own.tickDuration)
	// 计算任务应该放入的槽位置 (考虑当前指针位置)
	pos := (wheel.currentPosition + int(ticks)%wheel.slotNumber) % wheel.slotNumber
	// 将任务添加到对应的槽中
	wheel.slots[pos].PushBack(task)
	// 保存到任务映射
	own.taskMap.Store(task.id, task)
}

// CancelTask 取消任务（与单级实现相同）
func (own *MultiLevelTimeWheel) CancelTask(tid int64, fn func()) bool {
	if value, loaded := own.taskMap.LoadAndDelete(tid); loaded {
		value.(*Task).cancel = true
		if fn != nil {
			fn()
		}

		return true
	}

	return false
}

func (own *MultiLevelTimeWheel) callback(task *Task) {
	// 任务取消或重复执行次数结束
	if task.cancel || task.repeats == 0 {
		return
	}

	// 若 task.repeats>0 执行一次任务，则重复次数减1
	if task.repeats > 0 {
		task.repeats--
	}

	// 执行任务
	own.waitGroup.Add(1)
	go func(t *Task) {
		// 执行任务
		t.job()

		// 处理重复任务
		if t.period > 0 && t.repeats != 0 {
			t.delay = time.Now().Add(t.period).UnixNano()
			own.taskChan <- t // 重新添加
		} else {
			own.taskMap.Delete(t.id)
		}

		own.waitGroup.Done()
	}(task)
}
